//"Copyright [2021] <Copyright Shengkai Lin>"

#include <string>
#include <algorithm>
#include <memory>

#include "inet/common/packet/Packet.h"
#include "inet/networklayer/common/L3AddressResolver.h"
#include "inet/transportlayer/udp/Udp.h"
#include "DcqcnRoceConnection.h"
#include "RoceAdapter.h"
#include "RoceMsg_m.h"
#include <iostream>

namespace roce {

Define_Module(RoceAdapter);

void RoceAdapter::processStart() {
    socket_.setOutputGate(gate("socketOut"));
    roce_type_ = RoceType(static_cast<int>(par("roceType")));

    local_idx_ = par("localIdx");
    local_port_ = par("localPort");
    dst_port_ = par("destPort");
    dont_fragment_ = par("dontFragment");
    socket_.bind(GetL3Address(local_idx_), local_port_);
    setSocketOptions();
}

void RoceAdapter::HandleUpperLayerPacket(FlowMsg* f_msg) {
    Flow the_flow = f_msg->getThe_flow();
    std::unique_ptr<RoceConnectionBase>& roce_connection =
        RoceAdapter::FindOrCreateRoceConnection(the_flow.destIdx);
    // std::cout << "before HandlePacketFromHigherLayer" << std::endl;
    roce_connection->HandlePacketFromHigherLayer(*f_msg);
    // std::cout << "after HandlePacketFromHigherLayer" << std::endl;
    CheckAndSend(the_flow.destIdx);
    delete f_msg;
}

void RoceAdapter::HandleLowerLayerPacket(inet::Packet *packet) {
    const auto &payload = packet->peekAtFront<inet::RoceMsg>();
    Flow the_flow = payload->getThe_flow();
    int dst_idx = the_flow.srcIdx;
    std::unique_ptr<RoceConnectionBase>& roce_connection =
        RoceAdapter::FindOrCreateRoceConnection(dst_idx);
    roce_connection->HandlePacketFromLowerLayer(*packet);
    delete packet;
    inet::Packet* control_packet = roce_connection->GetNextControlPacket();
    while(control_packet) {
        EV << "Socket send Control packet: " << control_packet->getName()
           <<" data size:"<< control_packet->getTotalLength() << endl;
        socket_.sendTo(control_packet, GetL3Address(dst_idx), dst_port_);
        control_packet = roce_connection->GetNextControlPacket();
    }
    FlowMsg* f_msg = roce_connection->GetReceivedFlowMessage();
    if (f_msg) {
        send(f_msg, "gpu_port$o");
    }
}

void RoceAdapter::handleMessageWhenUp(omnetpp::cMessage *msg) {
    if (msg->arrivedOn("gpu_port$i")) {
        HandleUpperLayerPacket(dynamic_cast<FlowMsg *> (msg));
    } else if (msg->isSelfMessage()) {
        switch (msg->getKind()) {
            case START: {
                ASSERT(msg == selfMsg);
                // Start and stop messages are deleted elsewhere.
                processStart();
                break;
            }
            case SEND: {
                int dst = dynamic_cast<inet::CheckSendMsg*>(msg)->getDst();
                CheckAndSend(dst);
                delete msg;
                break;
            }
            case STOP: {
                ASSERT(msg == selfMsg);
                processStop();
                break;
            }
            default: {
                throw omnetpp::cRuntimeError(
                    "Invalid kind %d in self message", selfMsg->getKind());
            }
        }
    } else if (msg->arrivedOn("socketIn")) {
        HandleLowerLayerPacket(omnetpp::check_and_cast<inet::Packet *>(msg));
    } else {
        error("INET adapter have error!");
    }
}

inet::L3Address RoceAdapter::GetL3Address(int idx) {
    char local_address[15];
    snprintf(local_address, sizeof(local_address), "192.168.%d.%d",
        (idx + 1) / 256, (idx + 1) % 256);
    EV_INFO << "LocalIdx" << idx << endl;
    EV_INFO << "LocalAddress" << local_address << endl;
    inet::L3Address l3_address;
    inet::L3AddressResolver().tryParse(l3_address, local_address, 1);
    return l3_address;
}

std::unique_ptr<RoceConnectionBase>& RoceAdapter::FindOrCreateRoceConnection(
    int dst_idx) {
    if (dst_to_roce_connections_.find(dst_idx) ==
        dst_to_roce_connections_.end()) {
        RoceConnectionBase* roce_connection;
        switch (roce_type_) {
            case RoceType::DCQCN: {
                roce_connection = new DcqcnRoceConnection(
                    GetL3Address(local_idx_), GetL3Address(dst_idx),
                    local_port_, dst_port_, roce_type_);
                break;
            }
            default:
                error("Roce type not recognized");
        }
        dst_to_roce_connections_.emplace(
            dst_idx, std::unique_ptr<RoceConnectionBase>(roce_connection)); 
    }
    return dst_to_roce_connections_[dst_idx];
}

void RoceAdapter::CheckAndSend(int dst_idx){
    std::unique_ptr<RoceConnectionBase>& roce_connection =
        RoceAdapter::FindOrCreateRoceConnection(dst_idx);
    inet::Packet* packet;
    omnetpp::simtime_t intervalTime;
    // std::cout << "before GetNextDataPacket" << endl;
    packet = roce_connection->GetNextDataPacket(&intervalTime);
    // if (packet == nullptr) {
    //     std::cout << "CheckAndSend nullptr\n";
    // } 
    if (packet != nullptr) {
        // std::cout << packet->getName() << std::endl;
        //std::cout << "Socket send " << packet->getName() <<" data size:"<< packet->getTotalLength() << endl;
        emit(inet::packetSentSignal, packet);
        // std::cout << "Before socket send" << endl;
        socket_.sendTo(packet, GetL3Address(dst_idx), dst_port_);
        // std::cout << "Send data fragment in Roce " << endl;
        inet::CheckSendMsg *check_and_send_msg_ = new inet::CheckSendMsg();
        check_and_send_msg_->setKind(SEND);
        check_and_send_msg_->setDst(dst_idx);
        scheduleAt(omnetpp::simTime()+intervalTime, check_and_send_msg_);
    }
}

}  // namespace ai_sim
